import torch
from typing import Dict


def preprocess_batch(batch: Dict[str, torch.Tensor], device: torch.device = None, dtype: torch.dtype = None) -> Dict[str, torch.Tensor]:
    batch["img"] = batch["img"].to(device = device, dtype = dtype)
    return batch


def postprocess_batch(batch: Dict[str, torch.Tensor], device: torch.device = None, dtype: torch.dtype = None) -> Dict[str, torch.Tensor]:
    img: torch.Tensor = batch["img"][-1]
    ev_img = torch.zeros_like(img[:, :3])
    ev_img[:, 0] = (img[:, -2] * 8) + (img[:, 2] / 2)
    ev_img[:, 1] = (img[:, 1] / 2)
    ev_img[:, 2] = (img[:, -1] * 8) + (img[:, 0] / 2)
    batch["img"] = ev_img.to(device = device, dtype = dtype)
    return batch